{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Intralayer interatomic distance analysis of VASP AIMD trajectory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please cite DOI: 10.1021/acs.jpcc.9b04863 & DOI: 10.26434/chemrxiv.11811660.\n",
      "This script analyzes intralayer interatomic distances from VASP AIMD trajectory.\n",
      "This script is interactive and requires user input. Please read carefully before proceeding:\n",
      "\n",
      "Note 1: Intended only for periodic slab models with vacuum along the z-direction.\n",
      "Note 2: Intended only for VASP NVT simulations.\n",
      "Note 3: Requires PBC calculator from sitator package.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/davidlim/anaconda3/lib/python3.7/site-packages/tqdm/autonotebook/__init__.py:14: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  \" (e.g. in jupyter console)\", TqdmExperimentalWarning)\n"
     ]
    }
   ],
   "source": [
    "from sitator.util import PBCCalculator    # Use PBC calculator from sitator package\n",
    "import os\n",
    "import math\n",
    "import numpy as np\n",
    "import time\n",
    "from tqdm import tqdm\n",
    "\n",
    "pwd = os.getcwd()\n",
    "\n",
    "print('Please cite DOI: 10.1021/acs.jpcc.9b04863 & DOI: 10.26434/chemrxiv.11811660.')\n",
    "print('This script analyzes intralayer interatomic distances from VASP AIMD trajectory.')\n",
    "print('This script is interactive and requires user input. Please read carefully before proceeding:\\n')\n",
    "\n",
    "print('Note 1: Intended only for periodic slab models with vacuum along the z-direction.')\n",
    "print('Note 2: Intended only for VASP NVT simulations.')\n",
    "print('Note 3: Requires PBC calculator from sitator package.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load XDATCAR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter the name of XDATCAR file : XDATCAR_cat\n",
      "Loading trajectory...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 179000/179000 [18:09<00:00, 164.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# VASP XDATCAR\n",
    "name = input('Enter the name of XDATCAR file : ')\n",
    "print('Loading trajectory...')\n",
    "\n",
    "file = pwd + '/' + name\n",
    "xdatcar = open(file,'r')\n",
    "xdatcar_lines = xdatcar.readlines()\n",
    "xdatcar_lines = [line.split() for line in xdatcar_lines]\n",
    "\n",
    "# Lattice matrix\n",
    "lat = np.array([[xdatcar_lines[2][0], xdatcar_lines[2][1], xdatcar_lines[2][2]], \\\n",
    "                [xdatcar_lines[3][0], xdatcar_lines[3][1], xdatcar_lines[3][2]], \\\n",
    "                [xdatcar_lines[4][0], xdatcar_lines[4][1], xdatcar_lines[4][2]]]).astype(float)\n",
    "pbc = PBCCalculator(lat)    # Periodic boundary condition\n",
    "\n",
    "species = xdatcar_lines[5]    # List of element names\n",
    "N_type = np.array(xdatcar_lines[6]).astype(int)    # List of number of atoms of each element\n",
    "\n",
    "xyz = []\n",
    "N_frame = 0\n",
    "for line_index, line in enumerate(tqdm(xdatcar_lines[7:])):\n",
    "\n",
    "    if 'Direct' in line:\n",
    "        N_frame += 1\n",
    "        start = line_index+1\n",
    "                \n",
    "        if N_frame > 1:\n",
    "            xyz.append(frame)\n",
    "        \n",
    "        frame = []\n",
    "    \n",
    "    if line_index >= start:\n",
    "        coord = [line[0], line[1], line[2]]\n",
    "        frame.append(coord)\n",
    "    \n",
    "    if line_index == len(xdatcar_lines[7:])-1:\n",
    "        N_atom = len(frame)\n",
    "        xyz.append(frame)\n",
    "        xyz = np.array(xyz).astype(float)\n",
    "        xyz = np.dot(xyz, lat)    # Cartesian coordinates\n",
    "        \n",
    "print('Done.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract slab information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter the list of metal elements, separated by spaces : Ag Pd\n",
      "Enter the list of adsorbate elements, separated by spaces : H\n",
      "Enter the ID of a nearest-neighbor atom that lies in the same layer as atom #1 : 2\n",
      "Enter the ID of another nearest-neighbor atom that lies in the same layer as atom #1 : 8\n",
      "Enter the ID of an atom that lies one layer above atom #1 : 50\n",
      "Enter the number of layers in the slab, excluding the deposit/adsorbate layer(s) : 3\n",
      "Unit cell information & normal direction extracted.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "metals = input('Enter the list of metal elements, separated by spaces : ')\n",
    "metals = metals.split()\n",
    "ads = input('Enter the list of adsorbate elements, separated by spaces : ')\n",
    "ads = ads.split()\n",
    "\n",
    "N_metal = 0    # Number of metal atoms\n",
    "for s_index, s in enumerate(species):\n",
    "    if s in metals:\n",
    "        N_metal += N_type[s_index]\n",
    "N_ads = N_atom - N_metal    # Number of adsorbate atoms\n",
    "\n",
    "ID2 = int(input('Enter the ID of a nearest-neighbor atom that lies in the same layer as atom #1 : '))\n",
    "ID3 = int(input('Enter the ID of another nearest-neighbor atom that lies in the same layer as atom #1 : '))\n",
    "ID4 = int(input('Enter the ID of an atom that lies one layer above atom #1 : '))\n",
    "\n",
    "for index, coord in enumerate(xyz[0]):\n",
    "\n",
    "    if index+1 == 1:\n",
    "        atom1 = coord\n",
    "    \n",
    "    elif index+1 == ID2:\n",
    "        atom2 = coord\n",
    "    \n",
    "    elif index+1 == ID3:\n",
    "        atom3 = coord\n",
    "    \n",
    "    elif index+1 == ID4:\n",
    "        atom4 = coord\n",
    "\n",
    "a1 = np.subtract(atom2,atom1)    # 1st intralayer vector\n",
    "r12 = np.linalg.norm(a1)\n",
    "a1 = a1 / r12\n",
    "\n",
    "a2 = np.subtract(atom3,atom1)    # 2nd intralayer vector\n",
    "r13 = np.linalg.norm(a2)\n",
    "a2 = a2 / r13\n",
    "\n",
    "dr_avg = np.mean([r12, r13])    # In-plane NN distance\n",
    "\n",
    "a3 = np.cross(a1,a2)    # Normal vector\n",
    "if a3[2] < 0:\n",
    "    a3 = -a3    # Ensure normal vector along +z\n",
    "a3 = a3 / np.linalg.norm(a3)\n",
    "\n",
    "dn_avg = np.dot(atom4,a3) - np.dot(atom1,a3)    # Interlayer distance; n = normal component\n",
    "\n",
    "N_slab = int(input('Enter the number of layers in the slab, excluding the deposit/adsorbate layer(s) : '))\n",
    "        \n",
    "atom_next = 0\n",
    "z_layer = []    # Average layer heights\n",
    "while atom_next < N_metal-1:\n",
    "    z_bin = []\n",
    "    \n",
    "    for atom in range(atom_next,N_metal):\n",
    "        z0 = xyz[0][atom][2]    # Use last frame\n",
    "        z1 = xyz[0][atom+1][2]\n",
    "        dz = z1 - z0\n",
    "        \n",
    "        if abs(dz) > 0.20*dn_avg:\n",
    "            atom_next = atom+1\n",
    "            break\n",
    "            \n",
    "        elif z0 not in z_bin:\n",
    "            z_bin.append(z0)\n",
    "        \n",
    "        if atom == N_atom-2:\n",
    "            atom_next = atom+1\n",
    "            break\n",
    "        \n",
    "    z_bin = np.array(z_bin)\n",
    "    z_avg = np.mean(z_bin)\n",
    "    z_layer.append(z_avg)\n",
    "    \n",
    "N_layer = len(z_layer)\n",
    "\n",
    "print('Unit cell information & normal direction extracted.\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intralayer interatomic distance analysis (metals only)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter the simulation timestep in fs : 1.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1000/1000 [00:08<00:00, 113.46it/s]\n"
     ]
    }
   ],
   "source": [
    "file_ravg = pwd + '/ravg.txt'\n",
    "out_ravg = open(file_ravg,'w')\n",
    "\n",
    "header = 'time (fs)\\t'\n",
    "for l in range(0,N_layer):\n",
    "    name = 'L' + str(l)\n",
    "    header += name + '\\t'\n",
    "header += '\\n'\n",
    "out_ravg.write(header)\n",
    "\n",
    "dt = float(input('Enter the simulation timestep in fs : '))\n",
    "\n",
    "for s in tqdm(range(0,N_frame)):\n",
    "    t = s*dt\n",
    "    out = str(t) + '\\t'\n",
    "    \n",
    "    layer = {}    # Sort atoms by layer\n",
    "    for l in range(0,N_layer):\n",
    "        name = 'L' + str(l)\n",
    "        layer[name] = []\n",
    "    \n",
    "    for atom in range(0,N_metal):\n",
    "        coord = xyz[s][atom]\n",
    "        z = coord[2]\n",
    "        \n",
    "        for l in range(0,N_layer):\n",
    "            name = 'L' + str(l)\n",
    "            z_ref = z_layer[l]\n",
    "            dz = z - z_ref\n",
    "            if abs(dz) < 0.50*dn_avg:\n",
    "                layer[name].append(coord)\n",
    "                break\n",
    "    \n",
    "    for l in range(0,N_layer):\n",
    "        name = 'L' + str(l)\n",
    "        layer[name] = np.array(layer[name])\n",
    "        \n",
    "        N = layer[name].shape[0]\n",
    "        if N <= 1:\n",
    "            rij_avg = 0.0\n",
    "            out += str(rij_avg) + '\\t'\n",
    "            continue\n",
    "        \n",
    "        rij = pbc.pairwise_distances(layer[name])    # Pairwise distance matrix        \n",
    "        rij_bin = []\n",
    "        for i in range(0,N):\n",
    "            for j in range(0,N):\n",
    "                if (j > i) and (rij[i][j] < 1.20*dr_avg):\n",
    "                    rij_bin.append(rij[i][j])\n",
    "        rij_bin = np.array(rij_bin)\n",
    "\n",
    "        if rij_bin.shape[0] == 0:\n",
    "            rij_avg = 0.0\n",
    "        else:\n",
    "            rij_avg = np.mean(rij_bin)\n",
    "        \n",
    "        out += str(rij_avg) + '\\t'\n",
    "    \n",
    "    out += '\\n'\n",
    "    out_ravg.write(out)\n",
    "    \n",
    "out_ravg.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
